#!/usr/bin/env python3
# coding: utf-8
# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===========================================================================
import base64
import glob
import json
import logging
import os
import platform
import shlex
import shutil
import string
import subprocess
import sys
import time

sys.path.append(os.path.dirname(os.path.dirname(__file__)))

from utils import ROOT_PATH, Validator
from module_utils.path_manager import TmpPath
from module_utils.safe_file_handler import SafeFileHandler

try:
    from urllib.parse import urljoin
    from urllib.request import Request, urlopen

    PY3 = True
except ImportError:
    from urllib2 import Request, urlopen
    from urlparse import urljoin

    PY3 = False

LOG = logging.getLogger("ascend_deployer.scripts.nexus")

nexus_run_port = 58081 # default port for nexus
ACTIVE = "active"


def get_passwd():
    charset = string.digits + string.ascii_letters
    max_attempts = 100
    attempts = 0

    while attempts < max_attempts:
        attempts += 1
        passwd = []
        try:
            with open("/dev/random", "rb") as f:
                while len(passwd) < 16:
                    byte_char = f.read(1)
                    byte = ord(byte_char) if isinstance(byte_char, str) else byte_char[0]
                    if byte < 248:
                        index = byte % len(charset)
                        passwd.append(charset[index])
        except (OSError, IOError) as e:
            raise Exception("Failed to read from /dev/random") from e

        password_str = "".join(passwd)
        has_digit = any(c in string.digits for c in password_str)
        has_lower = any(c in string.ascii_lowercase for c in password_str)
        has_upper = any(c in string.ascii_uppercase for c in password_str)

        if sum([has_digit, has_lower, has_upper]) >= 2:
            return password_str

    raise RuntimeError("Failed to generate a valid password after 100 attempts")


class Command:
    """
    Command is a class mainly to execute the shell commands
    """
    @staticmethod
    def run(cmd, ignore_errors=False, log=True):
        """
        run is a method to execute the shell commands
        1. cmd: the shell command to be executed
        2. ignore_errors: whether to ignore errors during command execution
        3. log: whether to log the command execution process
        4. return: the stdout and stderr of the executed command
        """
        if log:
            LOG.info("nexus: {}".format(cmd).center(120, "-"))
        result = subprocess.Popen(
            shlex.split(cmd),
            shell=False,
            universal_newlines=True,
            stderr=subprocess.PIPE,
            stdout=subprocess.PIPE,
        )
        out, err = result.communicate()
        if result.returncode != 0 and not ignore_errors:
            raise RuntimeError("run cmd: {} failed, reason: {}".format(cmd, err))
        if log:
            for line in out.splitlines():
                if isinstance(line, bytes):
                    # Convert `line` to Unicode if it's a byte string
                    try:
                        line = line.decode('utf-8')
                    except UnicodeDecodeError:
                        LOG.error("Error: Unable to decode the byte string using UTF-8")
                        continue
                LOG.info(line)
        return out, err

class OsRepository:
    NEXUS_USER = "admin"
    nexus_passwd = get_passwd()
    gpg_passwd = get_passwd()

    def __init__(self, ip=None, port=58081):
        try:
            env_nexus_run_ip = ""
            ssh_connection = os.environ.get("SSH_CONNECTION", "")
            if ssh_connection:
                ip_fields = ssh_connection.split()
                if len(ip_fields) >= 3:
                    env_nexus_run_ip = ip_fields[2]
            self.nexus_run_ip = ip or env_nexus_run_ip
            if not Validator().is_valid_ip(self.nexus_run_ip):
                raise RuntimeError("nexus_run_ip is invalid, please check env variable SSH_CONNECTION")
            global nexus_run_port
            nexus_run_port = port
            self.working_on_ipv6 = False
            if ":" in self.nexus_run_ip:  # ipv6格式需要用括号包住域名部分
                self.nexus_run_ip = "[%s]" % self.nexus_run_ip
                self.working_on_ipv6 = True
        except KeyError:
            raise RuntimeError("Get environment variable SSH_CONNECTION failed,maybe switch users after SSH connection")
        self.nexus_url = "http://{}:{}".format(self.nexus_run_ip, nexus_run_port)
        self.config = "{}/scripts/nexus_config.json".format(ROOT_PATH)
        self.arch = platform.machine()
        os.environ.pop("http_proxy", "")
        os.environ.pop("https_proxy", "")
        # 修复：使用SafeFileHandler.safe_read方法而不是with语句
        self.config_content = SafeFileHandler.safe_read_json(self.config)
        os.chmod(self.config, 0o640)
        self.nexus_data_dir = os.path.join(TmpPath.ROOT, "nexus-data")
        self.nexus_image_name = self.config_content.get("image")
        nexus_dir = os.path.join(ROOT_PATH, "resources", "nexus")
        try:
            self.nexus_image = glob.glob("{}/nexus*{}.tar".format(nexus_dir, self.arch))[0]
        except IndexError:
            raise RuntimeError(
                "The nexus image does not exist. Ensure that the nexus image is in the {} directory.".format(nexus_dir)
            )
        auth = base64.b64encode("{}:{}".format(self.NEXUS_USER, self.nexus_passwd).encode()).decode()
        self.post_headers = {
            "Content-Type": "application/json",
            "Authorization": "Basic {}".format(auth),
        }
        self.upload_headers = {
            "Authorization": "Basic {}".format(auth),
            "Content-Type": "application/octet-stream",
        }

    @staticmethod
    def run_request(url, data, headers, method, timeout=20):
        if PY3:
            req = Request(url=url, data=data, headers=headers, method=method)
        else:
            req = Request(url=url, data=data, headers=headers)
            def _get_method():
                return method
            req.get_method = _get_method
        urlopen(req, timeout=timeout)

    @staticmethod
    def get_download_os_info():
        os_list_dir = "{}/resources".format(ROOT_PATH)
        os_list = [os_item for os_item in glob.glob("{}/*aarch64".format(os_list_dir))]
        os_list.extend(os_item for os_item in glob.glob("{}/*x86_64".format(os_list_dir)))
        return os_list
    
    @staticmethod
    def open_firewall_port():
        """
        Open firewall rules for Nexus port
        """
        try:
            # Check and open firewalld port
            out, _ = Command.run("systemctl is-active firewalld", ignore_errors=True, log=False)
            if ACTIVE in out:
                cmd = "firewall-cmd --permanent --add-port={}/tcp".format(nexus_run_port)
                Command.run(cmd, ignore_errors=True)
                Command.run("firewall-cmd --reload", ignore_errors=True)
                LOG.info("Opening port {} using firewalld".format(nexus_run_port))
                return

            # Check and open iptables port (without overwriting existing config)
            out, _ = Command.run("systemctl is-active iptables", ignore_errors=True, log=False)
            if ACTIVE in out:
                # 只添加规则，不保存到文件（规则在重启后会丢失）
                cmd = "iptables -I INPUT -p tcp --dport {} -j ACCEPT".format(nexus_run_port)
                Command.run(cmd, ignore_errors=True)
                LOG.info("Added iptables rule for port {} (not persisted across reboots)".format(nexus_run_port))
                return

            # Check and open ufw port
            out, _ = Command.run("ufw status", ignore_errors=True, log=False)
            if ACTIVE in out:
                cmd = "ufw allow {}".format(nexus_run_port)
                Command.run(cmd, ignore_errors=True)
                LOG.info("Opening port {} using ufw".format(nexus_run_port))
                return

            LOG.info("No active firewall service detected, skipping port opening")
        except Exception as e:
            LOG.warning("Error opening firewall port: {}".format(str(e)))

    @staticmethod
    def close_firewall_port():
        """
        Close previously opened firewall rules for Nexus port
        """
        try:
            # Check and close firewalld port
            out, _ = Command.run("systemctl is-active firewalld", ignore_errors=True, log=False)
            if ACTIVE in out:
                cmd = "firewall-cmd --permanent --remove-port={}/tcp".format(nexus_run_port)
                Command.run(cmd, ignore_errors=True)
                Command.run("firewall-cmd --reload", ignore_errors=True)
                LOG.info("Closing port {} using firewalld".format(nexus_run_port))
                return

            # Check and close iptables port
            out, _ = Command.run("systemctl is-active iptables", ignore_errors=True, log=False)
            if ACTIVE in out:
                # remove the INPUT rule
                cmd = "iptables -D INPUT -p tcp --dport {} -j ACCEPT".format(nexus_run_port)
                Command.run(cmd, ignore_errors=True)
                LOG.info("Removed iptables rule for port {}".format(nexus_run_port))
                return

            # Check and close ufw port
            out, _ = Command.run("ufw status", ignore_errors=True, log=False)
            if ACTIVE in out:
                cmd = "ufw delete allow {}".format(nexus_run_port)
                Command.run(cmd, ignore_errors=True)
                LOG.info("Closing port {} using ufw".format(nexus_run_port))
                return

            LOG.info("No active firewall service detected, skipping port closing")
        except Exception as e:
            LOG.warning("Error closing firewall port: {}".format(str(e)))

    @staticmethod
    def delete_nexus_container():
        out, err = Command.run("docker ps -a", ignore_errors=True)
        if "daemon running" in err:
            Command.run("systemctl daemon-reload")
            Command.run("systemctl restart docker")
            out, _ = Command.run("docker ps -a")
        if "nexus" in out.split():
            Command.run("docker rm -f nexus")
        Command.run("docker network rm ip6net_nexus", ignore_errors=True)

    @staticmethod
    def _check_nexus_status():
        timeout = 0
        while True:
            time.sleep(1)
            timeout += 1
            out, _ = Command.run("docker logs nexus", log=False)
            if "Nexus OSS" in out:
                break
            if timeout >= 300:
                raise RuntimeError("Nexus startup timeout")

    def init_nexus(self):
        self._run_nexus()
        self._check_nexus_status()
        tmp_passwd_file = "{}/admin.password".format(self.nexus_data_dir)
        old_passwd = SafeFileHandler.safe_read(tmp_passwd_file, "r").strip()
        auth = base64.b64encode("{}:{}".format(self.NEXUS_USER, old_passwd).encode()).decode()
        headers = {
            "accept": "application/json",
            "Content-Type": "text/plain",
            "Authorization": "Basic {}".format(auth),
        }
        url = "{}/service/rest/v1/security/users/admin/change-password".format(self.nexus_url)
        self.run_request(url, data=self.nexus_passwd.encode("utf-8"), headers=headers, method="PUT")

        url = "{}/service/rest/v1/security/anonymous".format(self.nexus_url)
        data = {
            "enabled": True,
            "userId": "anonymous",
            "realmName": "NexusAuthorizingRealm",
        }
        self.run_request(url, data=json.dumps(data).encode("utf-8"), headers=self.post_headers, method="PUT")

    def _get_os_files(self, os_dir):
        files = []
        for file_name in os.listdir(os_dir):
            file_path = os.path.join(os_dir, file_name)
            if os.path.isdir(file_path):
                files.extend(self._get_os_files(file_path))
            else:
                files.append(file_path)
        return files

    def _create_data_dir(self):
        if os.path.exists(self.nexus_data_dir):
            try:
                shutil.rmtree(self.nexus_data_dir)
            except OSError:
                Command.run("umount {}".format(self.nexus_data_dir))
                shutil.rmtree(self.nexus_data_dir)
        os.makedirs(self.nexus_data_dir, mode=0o700)
        Command.run("mount -t tmpfs tmpfs {}".format(self.nexus_data_dir))

    def _run_nexus(self):
        self.delete_nexus_container()
        self._create_data_dir()
        if self.working_on_ipv6:
            if not os.path.exists("/etc/docker/daemon.json"):
                os.makedirs("/etc/docker/", mode=0o755, exist_ok=True)
                SafeFileHandler.safe_write_json("/etc/docker/daemon.json", {}, indent=1)
            docker_settings = SafeFileHandler.safe_read_json("/etc/docker/daemon.json")
            docker_settings["experimental"] = True
            docker_settings["ip6tables"] = True
            SafeFileHandler.safe_write_json("/etc/docker/daemon.json", docker_settings, indent=1)
            Command.run("systemctl daemon-reload")
            Command.run("systemctl restart docker")

        Command.run("docker load -i {}".format(self.nexus_image))
        start_nexus_cmd = "docker run -d --name nexus --network host -v {}:/nexus-data {}".format(
            self.nexus_data_dir, self.nexus_image_name
        )
        Command.run(start_nexus_cmd)

        # 开放防火墙端口
        self.open_firewall_port()

        # 修改/nexus-data/etc/nexus.properties 绑定 IP 和端口
        self._update_nexus_properties()

    def _update_nexus_properties(self):
        """
        等待 /nexus-data/etc/nexus.properties 生成，并修改绑定 IP 和端口
        """
        nexus_properties_path = os.path.join(self.nexus_data_dir, "etc", "nexus.properties")

        # 等待 nexus.properties 生成（最多等待 120 秒）
        timeout = 120
        start_time = time.time()
        while not os.path.exists(nexus_properties_path):
            if time.time() - start_time > timeout:
                raise RuntimeError("Timeout waiting for nexus.properties, Nexus may not have initialized correctly")
            time.sleep(2)  # 每 2 秒检查一次

        # 读取原配置
        props = SafeFileHandler.safe_read(nexus_properties_path, "r").splitlines(True)

        # 更新或添加 application-host 和 application-port
        new_lines = []
        host_found = False
        port_found = False
        for line in props:
            stripped = line.strip()

            if stripped.startswith("application-host="):
                value = stripped.split("=", 1)[1].strip()
                if not value or value == "0.0.0.0":  # 按需判断是否替换
                    value = self.nexus_run_ip
                new_lines.append("application-host={}\n".format(value))
                host_found = True

            elif stripped.startswith("application-port="):
                value = stripped.split("=", 1)[1].strip()
                if not value or not value.isdigit():
                    value = "58081"
                new_lines.append("application-port={}\n".format(value))
                port_found = True

            else:
                new_lines.append(line)

        if not host_found:
            new_lines.append("application-host={}\n".format(self.nexus_run_ip))
        if not port_found:
            new_lines.append("application-port=58081\n")

        # 写回nexus.properties
        SafeFileHandler.safe_write(nexus_properties_path, ''.join(new_lines))

        # 重启 Nexus 容器让配置生效
        Command.run("docker restart nexus")

class YumRepository(OsRepository):
    def create_blob(self):
        url = "{}/service/rest/v1/blobstores/file".format(self.nexus_url)
        data = {"softQuota": None, "path": "/nexus-data/blobs/yum", "name": "yum"}
        self.run_request(url, data=json.dumps(data).encode("utf-8"), headers=self.post_headers, method="POST")

    def create_repository(self):
        url = "{}/service/rest/v1/repositories/yum/hosted".format(self.nexus_url)
        os_info = self.get_download_os_info()
        for i in os_info:
            repository_name = os.path.basename(i)
            if repository_name in self.config_content["rpm_os"]:
                data = {
                    "name": repository_name,
                    "online": True,
                    "storage": {
                        "blobStoreName": "yum",
                        "strictContentTypeValidation": True,
                        "writePolicy": "ALLOW",
                    },
                    "cleanup": None,
                    "component": {
                        "proprietaryComponents": False,
                    },
                    "yum": {"repodataDepth": 0, "deployPolicy": "STRICT"},
                }
                self.run_request(url, data=json.dumps(data).encode("utf-8"), headers=self.post_headers, method="POST")

    def upload_rpm(self):
        base_url = "{}/repository/".format(self.nexus_url)
        download_os_list = self.get_download_os_info()
        for download_os in download_os_list:
            os_name = os.path.basename(download_os)
            if os_name not in self.config_content["rpm_os"]:
                continue
            os_deps = self._get_os_files(download_os)
            for os_dep in os_deps:
                file_content = SafeFileHandler.safe_read(os_dep, "rb")
                url = urljoin(base_url, "{}/{}".format(os.path.basename(download_os), os.path.basename(os_dep)))
                self.run_request(url, data=file_content, headers=self.upload_headers, method="PUT")


class AptRepository(OsRepository):
    def generate_gpg_key(self):
        gpg_dir = os.path.expanduser("~/.gnupg")
        if os.path.exists(gpg_dir):
            shutil.rmtree(gpg_dir)
        base_cmd = "gpg --gen-key --batch"
        gpg_data = """Key-Type: RSA
        Key-Length: 3072
        Name-Real: nexus
        Expire-Date: 3d
        Passphrase: {}
        """.format(
            self.gpg_passwd
        )
        result = subprocess.Popen(
            shlex.split(base_cmd),
            shell=False,
            universal_newlines=True,
            stderr=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
        )
        result.communicate(input=gpg_data)
        if result.returncode != 0:
            raise RuntimeError("GPG key generation failed")

    def export_gpg_key(self):
        gpg_key_dir = os.path.dirname(self.nexus_image)
        gpg_pub_key = "{}/nexus_pub.asc".format(gpg_key_dir)
        gpg_pri_key = "{}/nexus_pri.asc".format(gpg_key_dir)
        if os.path.exists(gpg_pub_key):
            os.unlink(gpg_pub_key)
        if os.path.exists(gpg_pri_key):
            os.unlink(gpg_pri_key)
        export_public_key_cmd = "gpg -a -o {} --export nexus".format(gpg_pub_key)
        export_private_key_cmd = (
            "gpg --batch --pinentry-mode=loopback --yes --passphrase {} "
            "-a -o {} --export-secret-key nexus".format(self.gpg_passwd, gpg_pri_key)
        )
        centos_release = "/etc/centos-release"
        if os.path.exists(centos_release):
            export_private_key_cmd = "gpg -a -o {} --export-secret-key nexus".format(gpg_pri_key)
        Command.run(export_public_key_cmd)
        os.chmod(gpg_pub_key, 0o644)
        Command.run(export_private_key_cmd, log=False)
        os.chmod(gpg_pri_key, 0o600)
        gpg_pri_content = SafeFileHandler.safe_read(gpg_pri_key, "r")
        return gpg_pri_content

    def create_blob(self):
        url = "{}/service/rest/v1/blobstores/file".format(self.nexus_url)
        data = {"softQuota": None, "path": "/nexus-data/blobs/apt", "name": "apt"}
        self.run_request(url, data=json.dumps(data).encode("utf-8"), headers=self.post_headers, method="POST")

    def create_repository(self, keypair):
        url = "{}/service/rest/v1/repositories/apt/hosted".format(self.nexus_url)
        os_info = self.get_download_os_info()
        for i in os_info:
            repository_name = os.path.basename(i)
            if repository_name in self.config_content["deb_os"]:
                codename = self.config_content["codename"][repository_name]
                data = {
                    "name": repository_name,
                    "online": True,
                    "storage": {
                        "blobStoreName": "apt",
                        "strictContentTypeValidation": True,
                        "writePolicy": "ALLOW",
                    },
                    "cleanup": None,
                    "component": {
                        "proprietaryComponents": False,
                    },
                    "apt": {"distribution": codename},
                    "aptSigning": {"keypair": keypair, "passphrase": self.gpg_passwd},
                }
                self.run_request(url, data=json.dumps(data).encode("utf-8"), headers=self.post_headers, method="POST")

    def upload_deb(self):
        base_url = "{}/repository/".format(self.nexus_url)
        download_os_list = self.get_download_os_info()
        for download_os in download_os_list:
            if os.path.basename(download_os) not in self.config_content["deb_os"]:
                continue
            os_deps = self._get_os_files(download_os)
            for os_dep in os_deps:
                file_content = SafeFileHandler.safe_read(os_dep, "rb")
                url = urljoin(base_url, "{}/".format(os.path.basename(download_os)))
                self.run_request(url, data=file_content, headers=self.upload_headers, method="POST")


def main(ip=None, port=58081):
    yum_repository = YumRepository(ip, port)
    download_os_list = [os.path.basename(os_name) for os_name in yum_repository.get_download_os_info()]
    have_rpm = any(os_item in yum_repository.config_content["rpm_os"] for os_item in download_os_list)
    have_deb = any(os_item in yum_repository.config_content["deb_os"] for os_item in download_os_list)
    yum_repository.init_nexus()
    if have_rpm:
        yum_repository.create_blob()
        yum_repository.create_repository()
        yum_repository.upload_rpm()
    if have_deb:
        apt_repository = AptRepository(ip, port)
        apt_repository.create_blob()
        apt_repository.generate_gpg_key()
        gpg_pair_key = apt_repository.export_gpg_key()
        if gpg_pair_key == "":
            raise RuntimeError("The file content is empty")
        apt_repository.create_repository(gpg_pair_key)
        apt_repository.upload_deb()


if __name__ == "__main__":
    main()
